import os
import time
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
import base64

class CryptoEngine:
    """
    加密引擎实现类，提供256位以上的加密强度
    支持文本和二进制数据加密
    """
    
    def __init__(self, key_size=32):  # 默认256位密钥
        self.key_size = key_size  # 密钥大小(字节)
        self.iv_size = 16  # 初始化向量大小(字节)
        
    def generate_key(self):
        """生成随机密钥"""
        start_time = time.time()
        key = os.urandom(self.key_size)
        generation_time = (time.time() - start_time) * 1000  # 转换为毫秒
        return key, generation_time
    
    def derive_key(self, password, salt=None):
        """从密码派生密钥"""
        if salt is None:
            salt = os.urandom(16)
        
        start_time = time.time()
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=self.key_size,
            salt=salt,
            iterations=100000,
        )
        key = kdf.derive(password.encode() if isinstance(password, str) else password)
        generation_time = (time.time() - start_time) * 1000  # 转换为毫秒
        
        return key, salt, generation_time
    
    def encrypt(self, data, key):
        """
        加密数据
        :param data: 待加密数据，可以是文本或二进制数据
        :param key: 加密密钥
        :return: (encrypted_data, iv)元组
        """
        if isinstance(data, str):
            data = data.encode('utf-8')
        
        # 生成随机IV
        iv = os.urandom(self.iv_size)
        
        # 使用PKCS7填充
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(data) + padder.finalize()
        
        # 使用AES-256-CBC加密
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        
        return encrypted_data, iv
    
    def decrypt(self, encrypted_data, key, iv):
        """
        解密数据
        :param encrypted_data: 加密后的数据
        :param key: 解密密钥
        :param iv: 初始化向量
        :return: 解密后的数据
        """
        # 使用AES-256-CBC解密
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        decryptor = cipher.decryptor()
        padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
        
        # 去除PKCS7填充
        unpadder = padding.PKCS7(128).unpadder()
        data = unpadder.update(padded_data) + unpadder.finalize()
        
        return data
    
    def encrypt_file(self, input_file_path, output_file_path, key):
        """
        加密文件
        :param input_file_path: 输入文件路径
        :param output_file_path: 输出文件路径
        :param key: 加密密钥
        :return: IV和加密时间(ms)
        """
        with open(input_file_path, 'rb') as f:
            data = f.read()
        
        start_time = time.time()
        encrypted_data, iv = self.encrypt(data, key)
        encryption_time = (time.time() - start_time) * 1000  # 转换为毫秒
        
        with open(output_file_path, 'wb') as f:
            f.write(iv + encrypted_data)  # 将IV附加到加密数据前面
            
        return iv, encryption_time
    
    def decrypt_file(self, input_file_path, output_file_path, key):
        """
        解密文件
        :param input_file_path: 输入文件路径
        :param output_file_path: 输出文件路径
        :param key: 解密密钥
        :return: 解密时间(ms)
        """
        with open(input_file_path, 'rb') as f:
            data = f.read()
            
        iv = data[:self.iv_size]
        encrypted_data = data[self.iv_size:]
        
        start_time = time.time()
        decrypted_data = self.decrypt(encrypted_data, key, iv)
        decryption_time = (time.time() - start_time) * 1000  # 转换为毫秒
        
        with open(output_file_path, 'wb') as f:
            f.write(decrypted_data)
            
        return decryption_time
    
    @staticmethod
    def base64_encode(data):
        """Base64编码数据"""
        return base64.b64encode(data).decode('utf-8')
    
    @staticmethod
    def base64_decode(data):
        """Base64解码数据"""
        return base64.b64decode(data.encode('utf-8') if isinstance(data, str) else data) 